from PyQt5.QtWidgets import QGraphicsScene, QGraphicsPixmapItem
from PyQt5.QtGui import QPixmap, QPen, QBrush, QColor
from PyQt5.QtCore import Qt, QRectF, QPointF


class CalibrationScene(QGraphicsScene):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.image_item = None
        self.rect_items = []  # 存储已绘制的矩形
        self.current_rect = None  # 当前正在绘制的矩形
        self.start_point = None
        self.end_point = None
        self.is_drawing = False
        self.image_path = None

    def set_image(self, image_path):
        self.image_path = image_path
        self.clear()
        self.rect_items = []
        self.current_rect = None
        self.start_point = None
        self.end_point = None
        self.is_drawing = False

        # 加载图像
        pixmap = QPixmap(image_path)
        self.image_item = QGraphicsPixmapItem(pixmap)
        self.addItem(self.image_item)
        self.setSceneRect(0, 0, pixmap.width(), pixmap.height())

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.start_point = event.scenePos()
            self.is_drawing = True

    def mouseMoveEvent(self, event):
        if self.is_drawing:
            self.end_point = event.scenePos()
            if self.current_rect:
                self.removeItem(self.current_rect)

            # 创建矩形
            rect = QRectF(self.start_point, self.end_point)
            self.current_rect = self.addRect(
                rect,
                QPen(QColor(255, 0, 0), 2, Qt.SolidLine),
                QBrush(QColor(255, 0, 0, 50)),
            )

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.is_drawing:
            self.end_point = event.scenePos()
            self.is_drawing = False

            # 添加最终矩形到列表
            rect = QRectF(self.start_point, self.end_point)
            rect_item = self.addRect(
                rect,
                QPen(QColor(255, 0, 0), 2, Qt.SolidLine),
                QBrush(QColor(255, 0, 0, 50)),
            )
            self.rect_items.append(rect_item)
            self.current_rect = None

    def clear_all(self):
        """清除所有标记并重新加载原图"""
        if self.image_path:
            self.set_image(self.image_path)

    def get_normalized_boxes(self):
        """获取归一化后的边界框坐标"""
        boxes = []
        for rect_item in self.rect_items:
            rect = rect_item.rect()
            # 计算中心点和宽高
            x_center = (
                (rect.left() + rect.right()) / 2 / self.image_item.pixmap().width()
            )
            y_center = (
                (rect.top() + rect.bottom()) / 2 / self.image_item.pixmap().height()
            )
            width = rect.width() / self.image_item.pixmap().width()
            height = rect.height() / self.image_item.pixmap().height()

            # 确保坐标在合理范围内
            x_center = max(0, min(1, x_center))
            y_center = max(0, min(1, y_center))
            width = max(0, min(1, width))
            height = max(0, min(1, height))

            boxes.append((0, x_center, y_center, width, height))  # 0 是类别ID
        return boxes
